#!/usr/bin/env python
# coding: utf-8

# # Null-text inversion + Editing with Prompt-to-Prompt

# In[2]:


from typing import Optional, Union, Tuple, List, Callable, Dict
from tqdm.notebook import tqdm
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler
import torch.nn.functional as nnf
import numpy as np
import abc
import ptp_utils
import seq_aligner
import shutil
from torch.optim.adam import Adam
from PIL import Image


# For loading the Stable Diffusion using Diffusers, follow the instuctions https://huggingface.co/blog/stable_diffusion and update MY_TOKEN with your token.

# In[3]:


scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
MY_TOKEN = ''
LOW_RESOURCE = False 
NUM_DDIM_STEPS = 50
GUIDANCE_SCALE = 7.5
MAX_NUM_WORDS = 77
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
ldm_stable = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=MY_TOKEN, scheduler=scheduler).to(device)
try:
    ldm_stable.disable_xformers_memory_efficient_attention()
except AttributeError:
    print("Attribute disable_xformers_memory_efficient_attention() is missing")
tokenizer = ldm_stable.tokenizer

# -*- coding: utf-8 -*-
"""
Null-text inversion 讲解版（在用户给出的源码上只增加中文注释/文档，不改变逻辑）。

⚠️ 注意：本文件刻意“原汁原味”保留了源码中的实现与小瑕疵（见 load_512 的注释），
仅用注释解释其意图、形参/返回、与 DDIM 公式的对应关系，以及使用注意事项。

依赖：
- torch, numpy, PIL.Image
- diffusers 中的 UNet2DConditionModel + AutoencoderKL + DDIMScheduler
- 一个含有 .unet / .vae / .tokenizer / .scheduler 的 model 容器（ldm_stable）
- 全局常量：NUM_DDIM_STEPS, GUIDANCE_SCALE, device（源码中被引用）

核心流程：
1) invert():
   - init_prompt(prompt) 构造上下文 context = cat([uncond, cond])
   - 读入/裁剪/缩放图像（load_512） → image2latent 得到 x0 的 VAE latent
   - ddim_inversion(): 从 x0 开始，用 next_step 做 DDIM “前向加噪” 得到整条 latent 轨迹
   - null_optimization(): 逐步优化“无条件嵌入”（每个步长一份），令 prev_step(noise_pred) 精准重建该步目标 z_{t-1}
   - 返回：原图与重建图，最终 latent，以及每步优化得到的无条件嵌入列表

2) 数学对应（省略噪声项缩写）：
   - 估计 x0： x0 = (x_t - sqrt(1-α_t) * ε) / sqrt(α_t)
   - DDIM 上一步： x_{t-1} = sqrt(α_{t-1}) * x0 + sqrt(1-α_{t-1}) * ε
   - CFG 噪声： ε = ε_u + s * (ε_c - ε_u)

注：这里的 timestep 索引方向、以及 next_step/prev_step 的命名与常见实现略有“名不副实”的感觉，
但两者搭配使用时是自洽的；见每个函数内的详细注释。
"""

from typing import Union
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as nnf
from torch.optim import Adam
from tqdm import tqdm

# 假设以下常量在外部定义：
# NUM_DDIM_STEPS: int = 50
# GUIDANCE_SCALE: float = 7.5
# device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class NullInversion:
    """Null-text inversion 的最小实现：
    - DDIM 前向加噪生成反演轨迹
    - 逐步优化“无条件嵌入”，使得每一步 prev_step 能重建 DDIM 轨迹中的目标 z_{t-1}

    约定：self.model 需包含 .unet / .vae / .tokenizer / .scheduler
    """

    def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):
        """DDIM 的“上一时刻”一步（从 t → t-Δ），对应去噪方向。

        参数
        -----
        model_output : ε_t，UNet 预测的噪声（已做了 CFG 融合）
        timestep : 当前 t 索引（与 scheduler.timesteps 的索引一一对应）
        sample : x_t（当前 latent）

        返回
        -----
        prev_sample : x_{t-1}

        公式
        -----
        x0 = (x_t - sqrt(1-α_t) * ε) / sqrt(α_t)
        x_{t-1} = sqrt(α_{t-1}) * x0 + sqrt(1-α_{t-1}) * ε
        """
        # 计算“上一个”离散时刻的索引（按推理步长转化）
        prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps

        # 取 α 的累计乘积（alphas_cumprod）
        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
        alpha_prod_t_prev = self.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
        beta_prod_t = 1 - alpha_prod_t

        # 估计 x0
        pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
        # 方向项（与 ε 配比到 t-1）
        pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output

        # 合成 x_{t-1}
        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
        return prev_sample

    def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray]):
        """DDIM 的“下一时刻”一步（从 t-Δ → t），对应前向加噪方向。

        命名提醒：这里把 (timestep - Δ, next_timestep=timestep) 这样命，
        实际上是用 t 的 α 指标把 x_{t-Δ} 映射到 x_t。与上面的 prev_step 配套使用时是一致的。
        """
        # 令 `timestep` 变成 t-Δ，`next_timestep` 保持为当前 t（截顶到 999 避免越界）
        timestep, next_timestep = min(timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep

        # 取 α 的累计乘积（注意：timestep 现在是 t-Δ，next_timestep 是 t）
        alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]
        beta_prod_t = 1 - alpha_prod_t

        # 以 t-Δ 的 α 估计 x0
        next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
        # 方向项（映射到 t 的尺度）
        next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output

        # 合成 x_t
        next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
        return next_sample

    def get_noise_pred_single(self, latents, t, context):
        """给定单一路径的 encoder_hidden_states，返回 UNet 的噪声预测 ε。"""
        noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)["sample"]
        return noise_pred

    def get_noise_pred(self, latents, t, is_forward=True, context=None):
        """一口气完成：CFG 噪声融合 + （根据 is_forward）做一步 next_step 或 prev_step。

        - 当 is_forward=True：表示做“前向加噪”（用于 DDIM inversion 阶段）→ next_step
        - 当 is_forward=False：表示做“去噪方向”的一步（用于重建）→ prev_step
        """
        # 复制一份 batch=2 的输入，用于 uncond/cond 两路共享同一个 x_t
        latents_input = torch.cat([latents] * 2)

        if context is None:
            context = self.context  # [uncond, cond] 已经在 init_prompt 里准备

        # 采样阶段（is_forward=False）才使用 >1 的 CFG scale；生成反演轨迹时用 1 保持物理一致
        guidance_scale = 1 if is_forward else GUIDANCE_SCALE

        # 一次前向，拿到两路噪声预测
        noise_pred = self.model.unet(latents_input, t, encoder_hidden_states=context)["sample"]
        noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)

        # CFG 融合
        noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)

        # 根据方向选择一步
        if is_forward:
            latents = self.next_step(noise_pred, t, latents)
        else:
            latents = self.prev_step(noise_pred, t, latents)
        return latents

    @torch.no_grad()
    def init_prompt(self, prompt: str):
        """构造 context = cat([uncond_embeddings, cond_embeddings])。
        - uncond 使用空字符串 "" 作为负提示（标准 CFG 写法）。
        - cond 使用用户 prompt。
        """
        uncond_input = self.model.tokenizer(
            [""], padding="max_length", max_length=self.model.tokenizer.model_max_length,
            return_tensors="pt"
        )
        uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0]

        text_input = self.model.tokenizer(
            [prompt],
            padding="max_length",
            max_length=self.model.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0]

        self.context = torch.cat([uncond_embeddings, text_embeddings])
        self.prompt = prompt

    @torch.no_grad()
    def ddim_loop(self, latent):
        """从 x0 latent 开始，做 NUM_DDIM_STEPS 次“前向加噪”来生成 DDIM 反演轨迹。

        细节：
        - 这里取 t = timesteps[len-1-i]，若 scheduler.timesteps 是降序（常见），
          则该索引相当于从“小 t → 大 t”遍历，配合 next_step 的 t-Δ→t 映射，
          正好是把较“干净”的态逐步加噪到更“嘈杂”的态。
        - 内部只用有条件分支预测噪声（无 CFG），这在反演轨迹阶段常用于物理一致的加噪。
        """
        uncond_embeddings, cond_embeddings = self.context.chunk(2)
        all_latent = [latent]
        latent = latent.clone().detach()
        for i in range(NUM_DDIM_STEPS):
            t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1]
            noise_pred = self.get_noise_pred_single(latent, t, cond_embeddings)
            latent = self.next_step(noise_pred, t, latent)
            all_latent.append(latent)
        return all_latent

    @property
    def scheduler(self):
        return self.model.scheduler

    @torch.no_grad()
    def ddim_inversion(self, image):
        """图像 → latent(x0) → 解码重建一张图（检验对齐）→ 生成 DDIM 反演 latent 轨迹。
        返回 (重建图, 轨迹)。"""
        latent = self.image2latent(image)
        image_rec = self.latent2image(latent)
        ddim_latents = self.ddim_loop(latent)
        return image_rec, ddim_latents

    def null_optimization(self, latents, num_inner_steps, epsilon):
        """核心：逐步优化“无条件嵌入”，令每步 prev_step 重建上一时刻的目标 latent。

        参数
        -----
        latents : list[Tensor]
            由 ddim_loop 生成的轨迹（长度=NUM_DDIM_STEPS+1），latents[-1] 是“最嘈杂”的末端。
        num_inner_steps : int
            每个时间步里，对 uncond 嵌入做多少次梯度迭代。
        epsilon : float
            早停阈值基数；实际阈值 = epsilon + i * 2e-5。

        返回
        -----
        uncond_embeddings_list : list[Tensor]
            每个时间步一份专属的“无条件嵌入”（在该步上过拟合到原图）。
        """
        uncond_embeddings, cond_embeddings = self.context.chunk(2)
        uncond_embeddings_list = []

        # 从轨迹末端（最嘈杂）开始往回做优化
        latent_cur = latents[-1]
        bar = tqdm(total=num_inner_steps * NUM_DDIM_STEPS)

        for i in range(NUM_DDIM_STEPS):
            # 本步的“可学习”无条件嵌入，互不干扰（每步单独一份参数）
            uncond_embeddings = uncond_embeddings.clone().detach()
            uncond_embeddings.requires_grad = True

            # 学习率随步数线性衰减（经验做法）
            optimizer = Adam([uncond_embeddings], lr=1e-2 * (1. - i / 100.))

            # 目标上一时刻 latent，以及对应的时间索引 t （ddim）
            latent_prev = latents[len(latents) - i - 2]
            t = self.model.scheduler.timesteps[i]

            # 先缓存条件噪声（不参与梯度，节省内循环算力）
            with torch.no_grad():
                noise_pred_cond = self.get_noise_pred_single(latent_cur, t, cond_embeddings)

            # 在该时间步内迭代优化 uncond 嵌入
            for j in range(num_inner_steps):
                # 预测无条件噪声
                noise_pred_uncond = self.get_noise_pred_single(latent_cur, t, uncond_embeddings)

                # CFG 组合噪声（采样阶段用 GUIDANCE_SCALE）
                noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_cond - noise_pred_uncond)

                # 用 prev_step 反推，得到复原的 z_{t-1}
                latents_prev_rec = self.prev_step(noise_pred, t, latent_cur)

                # 使复原的 z_{t-1} 逼近 DDIM 轨迹中的“真” z_{t-1}
                loss = nnf.mse_loss(latents_prev_rec, latent_prev)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                loss_item = loss.item()
                bar.update()

                # 早停：步数越靠后（i 越大）阈值略宽松
                if loss_item < epsilon + i * 2e-5:
                    break

            # 进度条对齐（若提前收敛，补满本步的进度）
            for j in range(j + 1, num_inner_steps):
                bar.update()

            # 记录本步学到的“无条件嵌入”（只取 batch 中的 uncond 那一条）
            uncond_embeddings_list.append(uncond_embeddings[:1].detach())

            # 用“学好的无条件 + 原始条件”真正走一步，把 latent_cur 更新为下一次循环的起点
            with torch.no_grad():
                context = torch.cat([uncond_embeddings, cond_embeddings])
                latent_cur = self.get_noise_pred(latent_cur, t, False, context)

        bar.close()
        return uncond_embeddings_list

    def invert(self, image_path: str, prompt: str, offsets=(0,0,0,0), num_inner_steps=10, early_stop_epsilon=1e-5, verbose=False):
        """完整的 Null-text inversion 管线。

        返回：
        - (image_gt, image_rec): 原图与从 VAE latent 解码的重建图（用于 sanity check）
        - ddim_latents[-1]: 反演轨迹的末端 latent（最嘈杂）即噪声
        - uncond_embeddings: 每步的无条件嵌入列表（用于后续编辑）
        """
        self.init_prompt(prompt)

        # 读取/裁剪/缩放到 512
        image_gt = load_512(image_path, *offsets)

        if verbose:
            print("DDIM inversion...")
        image_rec, ddim_latents = self.ddim_inversion(image_gt)

        if verbose:
            print("Null-text optimization...")
        uncond_embeddings = self.null_optimization(ddim_latents, num_inner_steps, early_stop_epsilon)

        return (image_gt, image_rec), ddim_latents[-1], uncond_embeddings

    def __init__(self, model):
        # 这里新建了一个 DDIMScheduler，但实际上后面并没有把它赋给 model；
        # 紧接着又用 self.model.scheduler.set_timesteps(NUM_DDIM_STEPS)。
        # 因此 `scheduler = DDIMScheduler(...)` 这行对象未被使用，似为历史遗留。
        scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False,
                                  set_alpha_to_one=False)
        self.model = model
        self.tokenizer = self.model.tokenizer

        # 设置推理步数（影响 timesteps 的跨度和 prev/next_step 中的 Δ）
        self.model.scheduler.set_timesteps(NUM_DDIM_STEPS)

        self.prompt = None
        self.context = None


# 使用方式示例（假设 ldm_stable 是预加载的 Stable Diffusion 管线容器）
# null_inversion = NullInversion(ldm_stable)


# ## Infernce Code

# In[6]:


@torch.no_grad()
def text2image_ldm_stable(
    model,
    prompt:  List[str],
    controller,
    num_inference_steps: int = 50,
    guidance_scale: Optional[float] = 7.5,
    generator: Optional[torch.Generator] = None,
    latent: Optional[torch.FloatTensor] = None,
    uncond_embeddings=None,
    start_time=50,
    return_type='image'
):
    batch_size = len(prompt)
    ptp_utils.register_attention_control(model, controller)
    height = width = 512
    
    text_input = model.tokenizer(
        prompt,
        padding="max_length",
        max_length=model.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    text_embeddings = model.text_encoder(text_input.input_ids.to(model.device))[0]
    max_length = text_input.input_ids.shape[-1]
    if uncond_embeddings is None:
        uncond_input = model.tokenizer(
            [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
        )
        uncond_embeddings_ = model.text_encoder(uncond_input.input_ids.to(model.device))[0]
    else:
        uncond_embeddings_ = None

    latent, latents = ptp_utils.init_latent(latent, model, height, width, generator, batch_size)
    model.scheduler.set_timesteps(num_inference_steps)
    for i, t in enumerate(tqdm(model.scheduler.timesteps[-start_time:])):
        if uncond_embeddings_ is None:
            context = torch.cat([uncond_embeddings[i].expand(*text_embeddings.shape), text_embeddings])
        else:
            context = torch.cat([uncond_embeddings_, text_embeddings])
        latents = ptp_utils.diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False)
        
    if return_type == 'image':
        image = ptp_utils.latent2image(model.vae, latents)
    else:
        image = latents
    return image, latent



def run_and_display(prompts, controller, latent=None, run_baseline=False, generator=None, uncond_embeddings=None, verbose=True):
    if run_baseline:
        print("w.o. prompt-to-prompt")
        images, latent = run_and_display(prompts, EmptyControl(), latent=latent, run_baseline=False, generator=generator)
        print("with prompt-to-prompt")
    images, x_t = text2image_ldm_stable(ldm_stable, prompts, controller, latent=latent, num_inference_steps=NUM_DDIM_STEPS, guidance_scale=GUIDANCE_SCALE, generator=generator, uncond_embeddings=uncond_embeddings)
    if verbose:
        ptp_utils.view_images(images)
    return images, x_t


